import glob
import os
import json
import os.path
from tqdm import tqdm

import xml.sax
from shutil import copyfile
from xml import etree
import cv2
import numpy as np

from utils.plots import plot_one_box

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def mkdir(path):  # 创建文件
    if not os.path.exists(path):
        os.makedirs(path)


def load_all_cvat_label(label_path):  # 加载cvat2d图像的标签
    cvat_label = dict()
    qianzhui_name_lists = []
    with open(label_path, "r") as f:
        tree = ET.parse(label_path)
        root = tree.getroot()
        image_node_list = root.findall("image")  # root中找到所有image关键词  所有节点找到
        bar = tqdm(image_node_list)
        for idx, image_node in enumerate(bar):
            bar.set_description("load_all_cvat_label: ----- {}/{}".format(idx + 1, len(bar)))
            image_dict = image_node.attrib
            image_name = image_dict["name"].split("/")[2].split(".")[0]  # 这个是图像的名称
            qianzhui_name_lists.append(image_name)
            image_label_dict = dict()
            box_node_list = image_node.findall("box")
            # cvat_label_name_list = []
            for box_node in box_node_list:
                box_dict = box_node.attrib
                # cvat_label_name = box_dict["label"]
                # cvat_label_name_list.append(cvat_label_name)
                image_label_dict[box_dict["label"]] = {"xtl": box_dict["xtl"], "ytl": box_dict["ytl"],
                                                       "xbr": box_dict["xbr"], "ybr": box_dict["ybr"]}  # 字典套字典
                # 内层字典key是label的名称  val还是一个字典  key是 坐标值索引  val是值
            cvat_label[image_name] = image_label_dict  # 把标签的box信息字典放进字典里面去

        cvat_label['qianzhui'] = qianzhui_name_lists  # 信息名字前缀组成的list
        return cvat_label


def load_all_map_relation(map_path):  # 加载图像中label标签和3d点云中label标签对应关系的函数
    map_all_dict = dict()

    image_name = map_path.split("/")[-1].split('.')[0]
    with open(os.path.join(map_path), "r") as f:
        map_str_list = f.readlines()
        map_dict = dict()
        for map_str in map_str_list:
            lidar_name, cvat_name = map_str.split("'")[1], map_str.split("'")[3]  # 在3d标签中的名称与在2dcvat标签中的名称的一个映射关系
            map_dict[lidar_name] = cvat_name

    return map_dict


# 合并cvat和lidar标签
def concat_lidar_cvat(images_labels_path, threed_labels_path, map_relation_path, save_path):
    concat_label_path = save_path
    mkdir(concat_label_path)
    cvat_label_dict = load_all_cvat_label(images_labels_path)  # 加载并处理2d cvat的标签数据
    qianzhui_name_lists = cvat_label_dict['qianzhui']
    for qianzhui_name_list in qianzhui_name_lists:
        with open(os.path.join(threed_labels_path, qianzhui_name_list + '.txt'), 'r') as f:
            info = f.readlines()
            for i in range(len(info)):
                save_info = info[i].split(' ')
                label_info = save_info[0] + '_' + save_info[1]  # 3d的label标签名字
                map_dict = load_all_map_relation(os.path.join(map_relation_path, qianzhui_name_list + '.txt'))
                cvat_name = map_dict[label_info]
                image_label_dict = cvat_label_dict[qianzhui_name_list]
                if cvat_name in image_label_dict.keys():
                    with open(os.path.join(save_path, qianzhui_name_list + '.txt'), 'a') as fp:
                        save_info_list = []
                        cvat_n = cvat_name.split('_')[0]
                        if cvat_n == "Truck":
                            cvat_n = 'MiningTruck'
                        save_info_list.append(cvat_n)
                        image_label_dict = cvat_label_dict[qianzhui_name_list]
                        xtl = float(image_label_dict[cvat_name]['xtl'])
                        ytl = float(image_label_dict[cvat_name]['ytl'])
                        xbr = float(image_label_dict[cvat_name]['xbr'])
                        ybr = float(image_label_dict[cvat_name]['ybr'])
                        x = str((xbr+xtl)/(1920*2))
                        y = str((ybr+ytl)/(1080*2))
                        w = str((xbr-xtl)/1920)
                        h = str((ybr-ytl)/1080)
                        save_info_list.append(x)
                        save_info_list.append(y)
                        save_info_list.append(w)
                        save_info_list.append(h)
                        save_info_list.append(save_info[11])
                        s = ' '.join(save_info_list)
                        fp.write(s)
                        fp.write('\n')




#  算是写完了  但是有一个问题是  ：  这是以3d框标注为标准进行查找    其实我看了  3d labels的标签是都比2dcvat标签多的  但是如果没有一一对应的话 怎么办？还打印不打印？  不打印了 如果2d标签没有的话 直接不生成txt了  也解决了

def start(images_labels_path, threed_labels_path, map_relation_path, save_path):
    print("开始转换---------------------------->")
    concat_lidar_cvat(images_labels_path, threed_labels_path, map_relation_path, save_path)
    print("转换完成！")


images_labels_path = "/run/user/1000/gvfs/sftp:host=10.10.10.114,user=fuyu/home/fuyu/code/convert_tools/boleidun/2d_labels/right_front/annotations.xml"  # 图片的labels的地址
threed_labels_path = "/run/user/1000/gvfs/sftp:host=10.10.10.114,user=fuyu/home/fuyu/code/convert_tools/boleidun/labels"  # 3d的labels的地址
save_path = "/run/user/1000/gvfs/sftp:host=10.10.10.114,user=fuyu/home/fuyu/code/convert_tools/boleidun/threed_labels"  # 最后保存得到的地址
map_relation_path = '/run/user/1000/gvfs/sftp:host=10.10.10.114,user=fuyu/home/fuyu/code/convert_tools/boleidun_cvat/map_relation/right_front'  # 映射关系的地址们

start(images_labels_path=images_labels_path, threed_labels_path=threed_labels_path,
      map_relation_path=map_relation_path, save_path=save_path)
